#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/3/21 20:27
# @Author  : Seven
# @File    : ImageSimilarity.py
# @Software: PyCharm
# function : 图像相似度计算
import os
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
import tensorflow.contrib.slim as slim
import tensorflow.contrib.slim.nets as nets


# 创建一个session，用于运行后续的graph。
tf.reset_default_graph()
tf.logging.set_verbosity(tf.logging.ERROR)
sess = tf.InteractiveSession()
# 声明一个可训练的Tensor变量（variable），作为网络的输入数据
image = tf.Variable(tf.zeros((299, 299, 3)))


# 从slim中调用inception的网络定义。
def network(image, reuse):
    preprocessed = tf.multiply(tf.subtract(tf.expand_dims(image, 0), 0.5), 2.0)
    arg_scope = nets.inception.inception_v3_arg_scope(weight_decay=0.0)
    with slim.arg_scope(arg_scope):
        logits, end_points = nets.inception.inception_v3(
            preprocessed, 1001, is_training=False, reuse=reuse)
        logits = logits[:, 1:]  # ignore background class
        probs = tf.nn.softmax(logits)  # probabilities
    return logits, probs, end_points


# 获取输出
logits, probs, end_points = network(image, reuse=False)
# 获取预训练权重
data_dir = 'model/'
checkpoint_filename = os.path.join(data_dir, 'inception_v3.ckpt')
# 将预训练权重恢复到模型中。
restore_vars = [var for var in tf.global_variables() if var.name.startswith('InceptionV3/')]
saver = tf.train.Saver(restore_vars)
saver.restore(sess, os.path.join(data_dir, 'inception_v3.ckpt'))


# 获取特征
def get_feature(img, feature_layer_name):
    p, feature_values = sess.run([probs, end_points], feed_dict={image: img})
    return feature_values[feature_layer_name].squeeze()


# 加载图片
image_data = []
images = os.listdir('images')
for idx, img in enumerate(images):
    img_path = os.path.join('images', img)
    img = PIL.Image.open(img_path)
    img = img.resize((299, 299))
    image_data.append(img)
    # plt.subplot(1, 8, idx + 1)
    # plt.axis('off')
    # plt.imshow(img)
    # plt.title('images[{}]'.format(idx))
    # plt.show()

# 图片数据输入网络之前，需要做一下预处理，缩放到299x299然后为[0,1]之间的浮点数。
layer = 'PreLogits'  # 这个layer就是分类器前的最后一层.
features = []
for img in image_data:
    img = (np.asarray(img) / 255.0).astype(np.float32)
    feature = get_feature(img, layer)
    features.append(feature)
feature_vectors = np.stack(features)

# 最后得到的feature_vectors是一个矩阵，其中每一行代表了一张图片的feature，也就是特征向量， 每一列为特征向量的一个分量。
score_feature_shape = 0
try:
    print(feature_vectors.shape)
    assert feature_vectors.shape == (7, 2048), 'shape mismatch!'
    score_feature_shape = 10
except Exception as ex:
    print(ex)

# 计算欧式距离
distance_euclidean = np.sum(
    np.power(feature_vectors, 2), axis=1, keepdims=True) + np.sum(
        np.power(feature_vectors, 2), axis=1,
        keepdims=True).T - 2 * np.dot(feature_vectors, feature_vectors.T)
# 计算余弦距离
features_norm = feature_vectors / np.linalg.norm(
    feature_vectors, axis=1)[:, np.newaxis]
distance_cosin = np.dot(features_norm, features_norm.T)

plt.figure(figsize=(12, 21))

for idx_img, plt_img in enumerate(image_data):
    order_euclidean = np.argsort(distance_euclidean[idx_img])
    order_cosin = np.argsort(distance_cosin[idx_img])[::-1]

    plt.subplot(14, 8, idx_img * 8 * 2 + 1)
    plt.axis('off')
    plt.imshow(plt_img)
    plt.show()
    exit()
    for idx_sim, i in enumerate(order_euclidean):
        similay_img = image_data[i]
        plt.subplot(14, 8, idx_img * 8 * 2 + idx_sim + 2)
        plt.axis('off')
        plt.title('eucl-[{}]'.format(i))

        plt.imshow(similay_img)
        plt.show()

    for idx_sim, i in enumerate(order_cosin):
        similay_img = image_data[i]
        plt.subplot(14, 8, 8 + idx_img * 8 * 2 + idx_sim + 2)
        plt.axis('off')
        plt.title('cos-[{}]'.format(i))

        plt.imshow(similay_img)
        plt.show()
